#!/usr/bin/env python3
"""
工具调用演示
使用本地 Ollama 运行的 Qwen 模型，演示如何调用外部工具（如IP查询API）
"""

from util.llm import qwen_chat
import json
import requests
import sys
import time

# IP查询工具函数 - 使用PearAPI真实接口
def query_ip_info(ip_address):
    """
    查询IP地址的地理位置信息
    使用PearAPI接口: https://api.pearktrue.cn/api/ip/details/?ip={ip_address}
    
    参数:
        ip_address (str): 需要查询的IP地址，支持IPv4和IPv6
    
    返回:
        dict: 包含IP地址信息的字典
    """
    try:
        # 检查是否提供了IP地址
        if not ip_address:
            return {
                "code": 400,
                "msg": "缺少必要的参数: ip地址",
                "data": {}
            }
        
        # 调用PearAPI查询IP信息
        api_url = f"https://api.pearktrue.cn/api/ip/details/?ip={ip_address}"
        response = requests.get(api_url, timeout=10)
        result = response.json()
        return result
    except requests.exceptions.RequestException as e:
        return {
            "code": 500,
            "msg": f"网络请求错误: {str(e)}",
            "data": {}
        }
    except json.JSONDecodeError as e:
        return {
            "code": 500,
            "msg": f"响应解析错误: {str(e)}",
            "data": {}
        }
    except Exception as e:
        return {
            "code": 500,
            "msg": f"未知错误: {str(e)}",
            "data": {}
        }

# 工具定义
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "query_ip_info",
            "description": "查询IP地址的地理位置信息，包括归属州、国家、地区、城市、经纬度等详细信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "ip_address": {
                        "type": "string",
                        "description": "需要查询的IP地址，支持IPv4和IPv6，例如：8.8.8.8 或 2001:4860:4860::8888"
                    }
                },
                "required": ["ip_address"]
            },
            "returns": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "integer",
                        "description": "状态码"
                    },
                    "msg": {
                        "type": "string",
                        "description": "状态信息"
                    },
                    "ip": {
                        "type": "string",
                        "description": "查询的IP地址"
                    },
                    "data": {
                        "type": "object",
                        "properties": {
                            "continent": {
                                "type": "string",
                                "description": "归属州"
                            },
                            "continentCode": {
                                "type": "string",
                                "description": "归属州代码"
                            },
                            "country": {
                                "type": "string",
                                "description": "国家"
                            },
                            "countryCode": {
                                "type": "string",
                                "description": "国家代码"
                            },
                            "region": {
                                "type": "string",
                                "description": "所属区"
                            },
                            "regionCode": {
                                "type": "string",
                                "description": "所属区代码"
                            },
                            "city": {
                                "type": "string",
                                "description": "城市"
                            },
                            "districts": {
                                "type": "string",
                                "description": "区县城市"
                            },
                            "address": {
                                "type": "string",
                                "description": "地址信息"
                            },
                            "lat": {
                                "type": "number",
                                "description": "纬度"
                            },
                            "lon": {
                                "type": "number",
                                "description": "经度"
                            },
                            "timezone": {
                                "type": "string",
                                "description": "时区信息"
                            }
                        }
                    }
                }
            }
        }
    }
]

def handle_tool_call(tool_call):
    """
    处理工具调用
    """
    # 打印tool_call的完整结构用于调试
    print(f"[调试信息] tool_call对象: {tool_call}")
    
    function_name = tool_call.function.name
    arguments_str = tool_call.function.arguments
    
    try:
        arguments = json.loads(arguments_str)
    except json.JSONDecodeError:
        print(f"[调试信息] 无法解析参数: {arguments_str}")
        arguments = {}
    
    # 展示工具调用日志
    print(f"[工具调用日志] 调用工具: {function_name}")
    print(f"[工具调用日志] 参数: {arguments}")
    
    if function_name == "query_ip_info":
        ip_address = arguments.get("ip_address", "")
        result = query_ip_info(ip_address)
        print(f"[工具调用日志] 执行结果: {str(result)[:100]}...")  # 只显示前100个字符
        return {
            "name": function_name,
            "arguments": json.dumps(arguments),
            "result": json.dumps(result, ensure_ascii=False)
        }
    
    return None

def format_ip_info(data):
    """
    格式化IP信息以便更好地展示
    """
    if data.get("code") != 200:
        return f"查询失败: {data.get('msg', '未知错误')}"
    
    ip_data = data.get("data", {})
    info_parts = [
        f"IP地址: {data.get('ip', 'N/A')}",
        f"归属州: {ip_data.get('continent', 'N/A')} ({ip_data.get('continentCode', 'N/A')})",
        f"国家: {ip_data.get('country', 'N/A')} ({ip_data.get('countryCode', 'N/A')})",
        f"地区: {ip_data.get('region', 'N/A')} ({ip_data.get('regionCode', 'N/A')})",
        f"城市: {ip_data.get('city', 'N/A')}",
        f"区县: {ip_data.get('districts', 'N/A')}",
        f"详细地址: {ip_data.get('address', 'N/A')}",
        f"经纬度: {ip_data.get('lat', 'N/A')}, {ip_data.get('lon', 'N/A')}",
        f"时区: {ip_data.get('timezone', 'N/A')}"
    ]
    
    return "\n".join(info_parts)

def stream_output(text):
    """
    流式输出文本，模拟打字机效果
    """
    for char in text:
        sys.stdout.write(char)
        sys.stdout.flush()
        time.sleep(0.02)
    print()  # 换行

def main():
    print("工具调用演示 (单轮对话)")
    print("您可以询问类似'查询IP地址159.226.171.49的位置'这样的问题")
    print("输入 'quit' 或 'exit' 退出程序")
    print("-" * 50)
    
    # 设置系统提示，告诉模型可以使用工具
    system_prompt = """
你是一个可以调用工具的智能助手。当用户询问与IP地址相关的问题时，
你应该使用query_ip_info工具来获取准确的信息。
只在需要时调用工具，不要为每个请求都调用工具。
如果工具调用返回了信息，使用这些信息来回答用户的问题。
""".strip()

    while True:
        try:
            # 获取用户输入
            user_input = input("\n你: ").strip()
            
            # 检查退出命令
            if user_input.lower() in ['quit', 'exit']:
                stream_output("AI: 再见！")
                break
            
            # 跳过空输入
            if not user_input:
                continue
            
            response = qwen_chat(
                system_prompt=system_prompt,
                user_prompt=user_input,
                tools=TOOLS,
                stream=False
            )
            
            # 调试信息：打印完整的response
            print(f"[调试信息] 完整response: {response}")
            print(f"[调试信息] response类型: {type(response)}")
            
            # 检查是否需要工具调用
            tool_calls = getattr(response, 'tool_calls', None)
            print(f"[调试信息] tool_calls: {tool_calls}")
            
            if tool_calls:
                # 展示工具调用提取过程的日志
                print(f"[工具调用日志] 检测到 {len(tool_calls)} 个工具调用")
                
                # 处理工具调用
                tool_results = []
                for i, tool_call in enumerate(tool_calls):
                    print(f"[工具调用日志] 处理第 {i + 1} 个工具调用")
                    tool_result = handle_tool_call(tool_call)
                    if tool_result:
                        tool_results.append(tool_result)
                        # 格式化显示工具返回结果
                        try:
                            result_data = json.loads(tool_result['result'])
                            formatted_result = format_ip_info(result_data)
                            stream_output(f"[工具返回结果]:\n{formatted_result}")
                        except:
                            stream_output(f"[工具返回结果]: {tool_result['result']}")
                
                # 构造工具调用结果消息
                tool_messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_input},
                ]
                
                for result in tool_results:
                    tool_messages.append({
                        "role": "tool",
                        "name": result["name"],
                        "content": result["result"]
                    })
                
                final_response = qwen_chat(
                    messages=tool_messages,
                    stream=True,
                    tools=TOOLS
                )
                
                # 流式输出最终回答
                sys.stdout.write("AI: ")
                sys.stdout.flush()
                full_response = ""
                for chunk in final_response:
                    chunk_content = chunk.choices[0].delta.content or ""
                    full_response += chunk_content
                    sys.stdout.write(chunk_content)
                    sys.stdout.flush()
                    time.sleep(0.01)
                print()  # 换行
                print("-" * 50)  # 添加分隔线使输出更清晰
            else:
                print("[调试信息] 没有检测到工具调用")
                response_stream = qwen_chat(
                    system_prompt=system_prompt,
                    user_prompt=user_input,
                    stream=True
                )
                
                sys.stdout.write("AI: ")
                sys.stdout.flush()
                full_response = ""
                for chunk in response_stream:
                    chunk_content = chunk.choices[0].delta.content or ""
                    full_response += chunk_content
                    sys.stdout.write(chunk_content)
                    sys.stdout.flush()
                    time.sleep(0.01)
                print()  # 换行
                print("-" * 50)  # 添加分隔线使输出更清晰
            
            # 文件结尾
        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
            break
        except Exception as e:
            print(f"\n发生错误: {e}")

if __name__ == "__main__":
    main()
